import bs4
from langchain import hub
from langchain_community.document_loaders import WebBaseLoader
from langchain_community.embeddings import OpenAIEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain.embeddings import HuggingFaceEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter

from langchain_community.llms import Ollama


# RAG (Retrieval-Augmented Generation) is a technique that enhances the accuracy and reliability of generative AI models by fetching facts from external sources. It helps large language models respond authoritatively to user queries by providing specific citations or references, much like a court clerk would do in a courtroom.

def format_docs(docs):
    return "\\n\\n".join(doc.page_content for doc in docs)


def get_client():
    llm = Ollama(model="llama3")
    return llm


def load_docs():
    loader = WebBaseLoader(
        web_paths=("https://blogs.nvidia.com/blog/what-is-retrieval-augmented-generation/",),
        bs_kwargs=dict(
            parse_only=bs4.SoupStrainer(
                class_=("entry-content", "entry-header", "entry-title")
            )
        ),
    )
    docs = loader.load()
    return docs


def getRetriever(docs):
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    splits = text_splitter.split_documents(docs)

    embeddings = OpenAIEmbeddings(openai_api_base='https://api.chatanywhere.tech/v1')

    vectorstore = Chroma.from_documents(documents=splits, embedding=embeddings)
    # embedding=HuggingFaceEmbeddings())
    retriever = vectorstore.as_retriever()
    return retriever


def get_prompt():
    prompt = hub.pull("rlm/rag-prompt")
    return prompt


def get_chains():
    llm = get_client()
    prompt = get_prompt()
    docs = load_docs()
    retriever = getRetriever(docs)
    rag_chain = (
            {"context": retriever | format_docs, "question": RunnablePassthrough()}
            | prompt
            | llm
            | StrOutputParser()
    )

    out = rag_chain.invoke("Introducing RAG")
    return out
